from peewee import *
from datetime import datetime
import json
from .base import BaseModel


class SleepMonitoringRecord(BaseModel):
    """
    脑波会话数据表
    用于存储完整的睡眠监测会话数据，包括原始数据和分析结果
    """
    
    # 基础信息
    session_id = CharField(max_length=36, unique=True, help_text='会话唯一标识符')
    room_id = CharField(max_length=36, help_text='房间唯一标识符')
    start_time = DateTimeField(help_text='开始时间')
    end_time = DateTimeField(help_text='结束时间')
    total_seconds = IntegerField(help_text='总时长（秒）')
    
    # 摘要信息
    dominant_state = CharField(max_length=20, null=True, help_text='主要脑状态')
    avg_confidence = DecimalField(max_digits=3, decimal_places=2, null=True, help_text='平均置信度')
    deep_sleep_duration = IntegerField(null=True, help_text='深度睡眠时长（秒）')
    light_sleep_duration = IntegerField(null=True, help_text='浅度睡眠时长（秒）')
    awake_duration = IntegerField(null=True, help_text='清醒时长（秒）')
    description = TextField(null=True, help_text='治疗师描述')
    
    # 详细数据（JSON格式存储）
    data_states = TextField(null=True, help_text='状态时序数据 [[timestamp, state], [timestamp, state], ...]')
    data_confidence = TextField(null=True, help_text='置信度时序数据 [[timestamp, confidence], [timestamp, confidence], ...]')
    data_delta_ratio = TextField(null=True, help_text='Delta波比例时序数据 [[timestamp, ratio], [timestamp, ratio], ...]')
    data_theta_ratio = TextField(null=True, help_text='Theta波比例时序数据 [[timestamp, ratio], [timestamp, ratio], ...]')
    data_alpha_ratio = TextField(null=True, help_text='Alpha波比例时序数据 [[timestamp, ratio], [timestamp, ratio], ...]')
    data_beta_ratio = TextField(null=True, help_text='Beta波比例时序数据 [[timestamp, ratio], [timestamp, ratio], ...]')
    data_gamma_ratio = TextField(null=True, help_text='Gamma波比例时序数据 [[timestamp, ratio], [timestamp, ratio], ...]')
    
    # 原始脑波数据（JSON格式存储）
    data_raw_channel_1 = TextField(null=True, help_text='通道1原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_2 = TextField(null=True, help_text='通道2原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_3 = TextField(null=True, help_text='通道3原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_4 = TextField(null=True, help_text='通道4原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_5 = TextField(null=True, help_text='通道5原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_6 = TextField(null=True, help_text='通道6原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_7 = TextField(null=True, help_text='通道7原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    data_raw_channel_8 = TextField(null=True, help_text='通道8原始数据时序 [[timestamp, [data_array]], [timestamp, [data_array]], ...]')
    
    # 数据质量指标时序数据
    data_signal_quality = TextField(null=True, help_text='信号质量时序数据 [[timestamp, quality], [timestamp, quality], ...]')
    data_battery_level = TextField(null=True, help_text='电池电量时序数据 [[timestamp, level], [timestamp, level], ...]')
    
    class Meta:
        table_name = 'sleep_monitoring_records'
        indexes = (
            (('session_id',), False, 'idx_sleep_session_id'),
            (('start_time',), False, 'idx_sleep_start_time'),
            (('dominant_state',), False, 'idx_sleep_dominant_state'),
            (('created_at',), False, 'idx_sleep_created_at'),
        )
    
    def set_states_data(self, data: list):
        """设置状态时序数据"""
        self.data_states = json.dumps(data)
    
    def get_states_data(self) -> list:
        """获取状态时序数据"""
        if self.data_states:
            return json.loads(self.data_states)
        return []
    
    def set_confidence_data(self, data: list):
        """设置置信度时序数据"""
        self.data_confidence = json.dumps(data)
    
    def get_confidence_data(self) -> list:
        """获取置信度时序数据"""
        if self.data_confidence:
            return json.loads(self.data_confidence)
        return []
    
    def set_delta_ratio_data(self, data: list):
        """设置Delta波比例时序数据"""
        self.data_delta_ratio = json.dumps(data)
    
    def get_delta_ratio_data(self) -> list:
        """获取Delta波比例时序数据"""
        if self.data_delta_ratio:
            return json.loads(self.data_delta_ratio)
        return []
    
    def set_theta_ratio_data(self, data: list):
        """设置Theta波比例时序数据"""
        self.data_theta_ratio = json.dumps(data)
    
    def get_theta_ratio_data(self) -> list:
        """获取Theta波比例时序数据"""
        if self.data_theta_ratio:
            return json.loads(self.data_theta_ratio)
        return []
    
    def set_alpha_ratio_data(self, data: list):
        """设置Alpha波比例时序数据"""
        self.data_alpha_ratio = json.dumps(data)
    
    def get_alpha_ratio_data(self) -> list:
        """获取Alpha波比例时序数据"""
        if self.data_alpha_ratio:
            return json.loads(self.data_alpha_ratio)
        return []
    
    def set_beta_ratio_data(self, data: list):
        """设置Beta波比例时序数据"""
        self.data_beta_ratio = json.dumps(data)
    
    def get_beta_ratio_data(self) -> list:
        """获取Beta波比例时序数据"""
        if self.data_beta_ratio:
            return json.loads(self.data_beta_ratio)
        return []
    
    def set_gamma_ratio_data(self, data: list):
        """设置Gamma波比例时序数据"""
        self.data_gamma_ratio = json.dumps(data)
    
    def get_gamma_ratio_data(self) -> list:
        """获取Gamma波比例时序数据"""
        if self.data_gamma_ratio:
            return json.loads(self.data_gamma_ratio)
        return []
    
    # 原始数据设置和获取方法
    def set_raw_channel_data(self, channel: int, data: list):
        """设置指定通道的原始数据时序"""
        field_name = f'data_raw_channel_{channel}'
        if hasattr(self, field_name):
            setattr(self, field_name, json.dumps(data))
    
    def get_raw_channel_data(self, channel: int) -> list:
        """获取指定通道的原始数据时序"""
        field_name = f'data_raw_channel_{channel}'
        if hasattr(self, field_name):
            data = getattr(self, field_name)
            return json.loads(data) if data else []
        return []
    
    def set_signal_quality_data(self, data: list):
        """设置信号质量时序数据"""
        self.data_signal_quality = json.dumps(data)
    
    def get_signal_quality_data(self) -> list:
        """获取信号质量时序数据"""
        if self.data_signal_quality:
            return json.loads(self.data_signal_quality)
        return []
    
    def set_battery_level_data(self, data: list):
        """设置电池电量时序数据"""
        self.data_battery_level = json.dumps(data)
    
    def get_battery_level_data(self) -> list:
        """获取电池电量时序数据"""
        if self.data_battery_level:
            return json.loads(self.data_battery_level)
        return []
    
    def get_all_brainwave_data(self) -> dict:
        """获取所有脑波数据"""
        return {
            'states': self.get_states_data(),
            'confidence': self.get_confidence_data(),
            'delta_ratio': self.get_delta_ratio_data(),
            'theta_ratio': self.get_theta_ratio_data(),
            'alpha_ratio': self.get_alpha_ratio_data(),
            'beta_ratio': self.get_beta_ratio_data(),
            'gamma_ratio': self.get_gamma_ratio_data()
        }
    
    def calculate_summary(self):
        """计算摘要信息"""
        states_data = self.get_states_data()
        confidence_data = self.get_confidence_data()
        
        if states_data:
            # 计算主要状态
            state_counts = {}
            for timestamp, state in states_data:
                state_counts[state] = state_counts.get(state, 0) + 1
            self.dominant_state = max(state_counts, key=state_counts.get)
            
            # 计算各状态时长（假设每秒一个数据点）
            self.deep_sleep_duration = state_counts.get('deep_sleep', 0)
            self.light_sleep_duration = state_counts.get('light_sleep', 0)
            self.awake_duration = state_counts.get('awake', 0)
        
        if confidence_data:
            # 计算平均置信度
            confidences = [conf for timestamp, conf in confidence_data]
            self.avg_confidence = sum(confidences) / len(confidences)


class Room(BaseModel):
    """
    房间模型
    用于管理不同的监测房间
    """
    
    room_id = CharField(max_length=36, unique=True, help_text='房间唯一标识符')
    room_name = CharField(max_length=100, help_text='房间名称')
    description = TextField(null=True, help_text='房间描述')
    is_active = BooleanField(default=True, help_text='是否激活')
    
    class Meta:
        table_name = 'rooms'
        indexes = (
            (('room_id',), False, 'idx_room_room_id'),
            (('is_active',), False, 'idx_room_is_active'),
        )


# 注意：原始脑波数据现在存储在SleepMonitoringRecord模型中，不再使用独立的RawEEGData模型